#!/usr/bin/env python3

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc3 as pm

import h5py


with h5py.File("timing_samples.h5", 'r') as file_in:
    hit = file_in["hit"][:]
    nonhit = file_in["nonhit"][:]

x=hit["PulseTime"]

import scipy.integrate as si

def p(x, T=1):
    if x<=0:
        return 0
    return 1/(x**5*(np.exp(1/(T*x)-1)))

def logp(x, T=1):
    if x<=0:
        return -10000
    return -5*np.log(x)-np.log(np.exp(1/(T*x))-1)

def cdf(x=10, T=1, mu=0, a=0.01):
    if x<=0:
        return 0
    return si.quad(lambda x:p(x, T=T), a, x)

def logcdf(x=10, T=1, mu=0, a=0.01):
    if x<=0:
        return -100
    return np.log(si.quad(lambda x:p(x, T=T), a, x))

class Plank(pm.Continuous):
    def __init__(self, T, mu, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.T = T
        self.mu = mu

    def random(self, point, size):
        pass


    def logp(self, v, T, mu=0):
        T, mu = self.T, self.mu
        p = logp(v-mu)
        return p - np.log(cdf(T=T))

    def logcdf(self, v, T, mu=0):
        T, mu = self.T, self.mu
        return logcdf(v-mu, T=T) / logcdf(T=T)

distr = 'Plank'
with pm.Model():
    mu = pm.Uniform('mu', lower=5, upper=20)
    T = pm.Uniform('T', lower=0.01, upper=0.1)
    Y = Plank("pulse", T=T, mu=mu, observed=x)
    # pulse = pm.Uniform('pulse', lower=Y+5, upper=Y+15, observed=x)

    posterior = pm.sample(1000, cores=1)
    az.plot_trace(posterior)

# with pm.Model():
#     alpha = pm.Uniform('alpha', lower=2, upper=4)
#     beta = pm.Uniform('beta', lower=90, upper=150)
#     mu = pm.Uniform('mu', lower=7, upper=20)
#     Y = pm.InverseGamma("gamma", alpha=alpha, beta=beta, mu=mu, observed=x)
#     # pulse = pm.Uniform('pulse', lower=Y+5, upper=Y+15, observed=x)

#     posterior = pm.sample(1000, cores=1)
#     az.plot_trace(posterior)

plt.title(f'posterior probability for {distr}')
plt.show()
